# -*- coding:utf-8 -*-
from __future__ import division
from __future__ import print_function
import sys
from pyspark import SparkContext
from pyspark import SQLContext, HiveContext

from pyspark.sql.functions import udf, broadcast
from pyspark.sql.types import *
from pyspark.ml import Pipeline
from pyspark.ml.regression import GBTRegressor, GBTRegressionModel
from pyspark.ml.feature import VectorIndexer
from pyspark.ml.evaluation import RegressionEvaluator
from pyspark.ml.linalg import Vectors, VectorUDT


if __name__ == '__main__':

    model_file = sys.argv[1]

    sc = SparkContext(appName="Train age predict model ")
    sqlContext = HiveContext(sc)

    feature_cols = [
        "avg_monthly_consum_amt",
        "avg_monthly_order_cnt",
        "last_3d_order_cnt",
        "last_7d_order_cnt",
        "last_15d_order_cnt",
        "last_30d_order_cnt",
        "last_90d_order_cnt",
        "last_180d_order_cnt",
        "last_3d_order_amt",
        "last_7d_order_amt",
        "last_15d_order_amt",
        "last_30d_order_amt",
        "last_90d_order_amt",
        "last_180d_order_amt",
        "last_3d_order_unit_amt",
        "last_7d_order_unit_amt",
        "last_15d_order_unit_amt",
        "last_30d_order_unit_amt",
        "last_90d_order_unit_amt",
        "last_180d_order_unit_amt",
        "is_last_3d_order",
        "is_last_7d_order",
        "is_last_15d_order",
        "is_last_30d_order",
        "is_last_90d_order",
        "is_last_180d_order",
        "last_3d_goods_cnt",
        "last_7d_goods_cnt",
        "last_15d_goods_cnt",
        "last_30d_goods_cnt",
        "last_90d_goods_cnt",
        "last_180d_goods_cnt",
        "last_3d_team_cnt",
        "last_7d_team_cnt",
        "last_15d_team_cnt",
        "last_30d_team_cnt",
        "last_90d_team_cnt",
        "last_180d_team_cnt",
        "last_3d_class1_cnt",
        "last_7d_class1_cnt",
        "last_15d_class1_cnt",
        "last_30d_class1_cnt",
        "last_90d_class1_cnt",
        "last_180d_class1_cnt",
        "last_3d_avg_daily_order_cnt",
        "last_7d_avg_daily_order_cnt",
        "last_15d_avg_daily_order_cnt",
        "last_30d_avg_daily_order_cnt",
        "last_90d_avg_daily_order_cnt",
        "last_180d_avg_daily_order_cnt",
        # "top3_ai_class1_list",
        "dndq_order_amt",
        "dndq_order_cnt",
        "dndq_goods_cnt",
        "dndq_team_cnt",
        "hfcz_order_amt",
        "hfcz_order_cnt",
        "hfcz_goods_cnt",
        "hfcz_team_cnt",
        "hnwy_order_amt",
        "hnwy_order_cnt",
        "hnwy_goods_cnt",
        "hnwy_team_cnt",
        "jjjz_order_amt",
        "jjjz_order_cnt",
        "jjjz_goods_cnt",
        "jjjz_team_cnt",
        "myyp_order_amt",
        "myyp_order_cnt",
        "myyp_goods_cnt",
        "myyp_team_cnt",
        "nnfs_order_amt",
        "nnfs_order_cnt",
        "nnfs_goods_cnt",
        "nnfs_team_cnt",
        "qt_order_amt",
        "qt_order_cnt",
        "qt_goods_cnt",
        "qt_team_cnt",
        "qcyp_order_amt",
        "qcyp_order_cnt",
        "qcyp_goods_cnt",
        "qcyp_team_cnt",
        "rybh_order_amt",
        "rybh_order_cnt",
        "rybh_goods_cnt",
        "rybh_team_cnt",
        "shfw_order_amt",
        "shfw_order_cnt",
        "shfw_goods_cnt",
        "shfw_team_cnt",
        "spcy_order_amt",
        "spcy_order_cnt",
        "spcy_goods_cnt",
        "spcy_team_cnt",
        "sjsm_order_amt",
        "sjsm_order_cnt",
        "sjsm_goods_cnt",
        "sjsm_team_cnt",
        "tsyx_order_amt",
        "tsyx_order_cnt",
        "tsyx_goods_cnt",
        "tsyx_team_cnt",
        "xlxb_order_amt",
        "xlxb_order_cnt",
        "xlxb_goods_cnt",
        "xlxb_team_cnt",
        "ydhw_order_amt",
        "ydhw_order_cnt",
        "ydhw_goods_cnt",
        "ydhw_team_cnt",
        "zbps_order_amt",
        "zbps_order_cnt",
        "zbps_goods_cnt",
        "zbps_team_cnt",
    ]

    udf_vectorize = udf(lambda *xs: Vectors.dense([float(x) for x in xs]), VectorUDT())
    udf_to_float = udf(lambda x: float(x), FloatType())

    sql = """
    select * from dw.dws_scrm_persona_customer_dim as a join dm_scrm.dm_persona_order_goods_index as b on a.yz_uid = b.buyer_id where a.age > 0 and a.age < 90
    """
    df = sqlContext.sql(sql)

    df = df.withColumn('features', udf_vectorize(*feature_cols)).withColumn('target', udf_to_float('age'))

    gbr = GBTRegressor(featuresCol='features', labelCol='target', maxIter=100, maxBins=50, maxDepth=10)

    train_df, test_df = df.randomSplit([0.8, 0.2])

    model = gbr.fit(train_df)

    predictions = model.transform(test_df)
    pred = predictions.select("prediction", 'age').take(100)

    print("预测\t标签")
    for d in pred:
        p = d['prediction']
        t = d['age']
        print('\t'.join([str(k) for k in ([p, t])]))

    # evaluator = (labelCol='final_read_num', predictionCol='prediction', metricName='rmse')
    evaluator = RegressionEvaluator(labelCol='target', predictionCol='prediction', metricName='rmse')
    rmse = evaluator.evaluate(predictions)
    print("RMSE on spark data is %g" % rmse)
    model.save(model_file)
    print(model)

    sc.stop()

